import argparse
import os
import sys
from pathlib import Path

import torch
from torchvision.utils import save_image

from gan.model import create_gan

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative


def main(opt):
    model, _, _ = create_gan(dict(type="Pretrained", checkpoint=opt.checkpoint))
    model.init_weights()

    z = torch.randn(64, 100)
    img = model(z)

    save_image(img, opt.result, normalize=True)


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--checkpoint",
        default=ROOT / "runs/gan/epoch_50.pth",
        help="the checkpoint file",
    )
    parser.add_argument(
        "--result", default=ROOT / "runs/gan/result.png", help="the result image file"
    )
    opt = parser.parse_args()
    return opt


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
